package com.aiwiown.snackmq.api.impl;

import com.aiwiown.snackmq.api.BatchMessageListener;
import com.aiwiown.snackmq.api.Consumer;
import com.aiwiown.snackmq.api.MessageListener;
import com.aiwiown.snackmq.api.OrderedMessageListener;
import com.aiwiown.snackmq.api.PullResult;
import com.aiwiown.snackmq.api.config.ConsumerConfig;
import com.aiwiown.snackmq.api.exception.SnackMQClientException;
import com.aiwiown.snackmq.common.dto.TopicInfo;
import com.aiwiown.snackmq.common.message.Message;
import com.aiwiown.snackmq.common.message.MessageStatus;
import com.aiwiown.snackmq.common.message.MessageType;
import com.aiwiown.snackmq.network.client.NettyClient;
import lombok.extern.slf4j.Slf4j;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 【已重构】"拉"模式消费者的默认实现。
 * 此版本通过组合 ConsumerDelegate 来实现共享逻辑，遵循“组合优于继承”原则。
 */
@Slf4j
public class DefaultPullConsumer implements Consumer {

    // --- 新增：委托类，封装共享逻辑 ---
    private final ConsumerDelegate delegate;

    // --- Pull模式特有的组件 ---
    private final Map<String, AtomicInteger> partitionCursors = new ConcurrentHashMap<>();
    private static final int MAX_COMMIT_REDIRECT_RETRIES = 5;


    /**
     * 【修改】此公共构造函数现在委托给新的包级私有构造函数，以支持测试。
     */
    public DefaultPullConsumer(String brokerAddress, ConsumerConfig consumerConfig) {
        this(new ConsumerDelegate(brokerAddress, consumerConfig));
    }

    /**
     * 【新增】包级私有的构造函数，用于在单元测试中注入一个 mock/spy 的 delegate。
     */
    DefaultPullConsumer(ConsumerDelegate delegate) {
        this.delegate = delegate;
    }

    @Override
    public void start() throws SnackMQClientException {
        delegate.start();
    }

    @Override
    public void close() throws IOException {
        delegate.close();
    }

    @Override
    public List<Message> pullList(String topic, int maxMessages, long timeoutMs) throws SnackMQClientException {
        PullResult result = pullResult(topic, maxMessages, timeoutMs);
        if (result != null && result.getMessages() != null && !result.getMessages().isEmpty()) {
            commitOffset(result.getMessages().get(result.getMessages().size() - 1));
            return result.getMessages();
        }
        return Collections.emptyList();
    }

    @Override
    public List<Message> pullList(String topic, int partition, long offset, int maxMessages, long timeoutMs) throws SnackMQClientException {
        PullResult result = pullResult(topic, partition, offset, maxMessages, timeoutMs);
        if (result != null && result.getMessages() != null && !result.getMessages().isEmpty()) {
            return result.getMessages();
        }
        return Collections.emptyList();
    }

    @Override
    public PullResult pullResult(String topic, int maxMessages, long timeoutMs) throws SnackMQClientException {
        if (!delegate.isStarted()) {
            throw new SnackMQClientException("Consumer has not been started. Please call start() before pulling messages.");
        }

        try {
            TopicInfo topicInfo = delegate.fetchTopicMetadata(topic).get(timeoutMs, TimeUnit.MILLISECONDS);
            if (topicInfo == null) {
                throw new SnackMQClientException("Failed to fetch metadata for topic: " + topic);
            }

            List<Integer> partitionIds = new ArrayList<>(topicInfo.getPartitionAssignments().keySet());
            if (partitionIds.isEmpty()) {
                log.warn("No partitions found for topic: {}", topic);
                return new PullResult(Collections.emptyList(), -1);
            }
            Collections.sort(partitionIds);
            long lastOffsetPulled = -1;

            long deadline = System.currentTimeMillis() + timeoutMs;
            for (int i = 0; i < partitionIds.size(); i++) {
                if (System.currentTimeMillis() >= deadline) break;

                AtomicInteger cursor = partitionCursors.computeIfAbsent(topic, k -> new AtomicInteger(0));
                int partitionId = partitionIds.get(cursor.getAndIncrement() % partitionIds.size());
                TopicPartition tp = new TopicPartition(topic, partitionId);

                long remainingTimeout = Math.max(0, deadline - System.currentTimeMillis());
                long lastCommittedOffset = fetchOffset(tp).get(remainingTimeout, TimeUnit.MILLISECONDS);
                long offsetToPull = lastCommittedOffset + 1;
                lastOffsetPulled = offsetToPull;

                Message pullRequest = Message.builder()
                        .type(MessageType.PULL_REQUEST)
                        .topic(topic)
                        .property("consumerGroup", this.delegate.getConsumerConfig().getConsumerGroup())
                        .property("partition", String.valueOf(partitionId))
                        .property("offset", String.valueOf(offsetToPull))
                        .property("maxMessages", String.valueOf(maxMessages))
                        .build();

                remainingTimeout = Math.max(0, deadline - System.currentTimeMillis());
                Message response = delegate.sendRequestToPartitionLeader(pullRequest, tp).get(remainingTimeout, TimeUnit.MILLISECONDS);

                if (response.getStatus() != MessageStatus.SUCCESS && response.getStatus() != MessageStatus.EMPTY) {
                    throw new SnackMQClientException("Pull request for partition " + tp + " failed: " + response.getErrorMessage());
                }

                List<Message> messages = response.getMessages();
                if (messages != null && !messages.isEmpty()) {
                    long nextFetchOffset = Long.parseLong(messages.get(messages.size() - 1).getProperty("offset")) + 1;
                    return new PullResult(messages, nextFetchOffset);
                }
            }

            log.debug("No new messages found after polling all partitions. Next fetch offset will be {}.", lastOffsetPulled);
            return new PullResult(Collections.emptyList(), lastOffsetPulled);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new SnackMQClientException("Pull operation was interrupted.", e);
        } catch (TimeoutException e) {
            return new PullResult(Collections.emptyList(), -1);
        } catch (ExecutionException e) {
            if (e.getCause() instanceof SnackMQClientException) {
                throw (SnackMQClientException) e.getCause();
            }
            throw new SnackMQClientException("Pull operation failed.", e.getCause());
        }
    }

    @Override
    public PullResult pullResult(String topic, int partition, long offset, int maxMessages, long timeoutMs) throws SnackMQClientException {
        if (!delegate.isStarted()) {
            throw new SnackMQClientException("Consumer has not been started. Please call start() before pulling messages.");
        }
        final long offsetToPull = offset + 1;
        try {
            TopicPartition tp = new TopicPartition(topic, partition);
            Message pullRequest = Message.builder()
                    .type(MessageType.PULL_REQUEST)
                    .topic(topic)
                    .property("consumerGroup", this.delegate.getConsumerConfig().getConsumerGroup())
                    .property("partition", String.valueOf(partition))
                    .property("offset", String.valueOf(offsetToPull))
                    .property("maxMessages", String.valueOf(maxMessages))
                    .build();

            Message response = delegate.sendRequestToPartitionLeader(pullRequest, tp).get(timeoutMs, TimeUnit.MILLISECONDS);

            if (response.getStatus() != MessageStatus.SUCCESS && response.getStatus() != MessageStatus.EMPTY) {
                throw new SnackMQClientException("Pull request for partition " + tp + " failed: " + response.getErrorMessage());
            }

            List<Message> messages = response.getMessages();
            if (messages != null && !messages.isEmpty()) {
                long nextFetchOffset = Long.parseLong(messages.get(messages.size() - 1).getProperty("offset")) + 1;
                return new PullResult(messages, nextFetchOffset);
            }

            return new PullResult(Collections.emptyList(), offsetToPull);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new SnackMQClientException("Pull operation was interrupted.", e);
        } catch (TimeoutException e) {
            log.warn("Pull operation for topic {} partition {} from offset {} timed out. Suggesting retry from the same offset.", topic, partition, offsetToPull);
            return new PullResult(Collections.emptyList(), offsetToPull);
        } catch (ExecutionException e) {
            if (e.getCause() instanceof SnackMQClientException) {
                throw (SnackMQClientException) e.getCause();
            }
            throw new SnackMQClientException("Pull operation failed.", e.getCause());
        }
    }

    @Override
    public void commitOffset(Message message) throws SnackMQClientException {
        if (!delegate.isStarted()) {
            throw new SnackMQClientException("Consumer has not been started.");
        }
        if (message == null || message.getProperty("partition") == null || message.getProperty("offset") == null) {
            throw new IllegalArgumentException("Message to commit must not be null and must contain 'partition' and 'offset' properties.");
        }
        Message commitRequest = Message.builder()
                .type(MessageType.COMMIT_OFFSET)
                .topic(message.getTopic())
                .property("consumerGroup", this.delegate.getConsumerConfig().getConsumerGroup())
                .property("partition", message.getProperty("partition"))
                .property("offset", message.getProperty("offset"))
                .build();

        // Re-implement the synchronous redirect logic using the async delegate
        int retries = MAX_COMMIT_REDIRECT_RETRIES;
        CompletableFuture<Message> future = delegate.sendRequestToMetadataLeader(commitRequest);

        while (retries-- > 0) {
            try {
                Message response = future.get(30, TimeUnit.SECONDS);
                if (response.getStatus() == MessageStatus.SUCCESS) {
                    return; // Success
                }
                if (response.getStatus() == MessageStatus.REDIRECT && response.getProperty("redirectAddress") != null) {
                    String newLeaderAddress = response.getProperty("redirectAddress");
                    log.warn("Commit offset redirected to {}. Retrying...", newLeaderAddress);
                    NettyClient newClient = delegate.getOrCreateClient(newLeaderAddress);
                    future = newClient.sendRequest(delegate.addAuthToken(commitRequest));
                    continue; // Retry the loop with the new future
                }
                throw new SnackMQClientException("Failed to commit offset: " + response.getErrorMessage());
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
                throw new SnackMQClientException("Commit offset operation was interrupted.", e);
            } catch (TimeoutException e) {
                throw new SnackMQClientException("Commit offset operation timed out.", e);
            } catch (ExecutionException e) {
                Throwable cause = e.getCause() != null ? e.getCause() : e;
                if (cause instanceof SnackMQClientException) throw (SnackMQClientException) cause;
                throw new SnackMQClientException("Commit offset operation failed.", cause);
            }
        }
        throw new SnackMQClientException("Failed to commit offset after " + MAX_COMMIT_REDIRECT_RETRIES + " retries.");
    }

    private CompletableFuture<Long> fetchOffset(TopicPartition tp) {
        Message request = Message.builder()
                .type(MessageType.FETCH_OFFSET)
                .topic(tp.getTopic())
                .property("consumerGroup", delegate.getConsumerConfig().getConsumerGroup())
                .property("partition", String.valueOf(tp.getPartitionId()))
                .build();

        return delegate.sendRequestToPartitionLeader(request, tp).thenApply(response -> {
            if (response.getStatus() == MessageStatus.SUCCESS && response.getProperty("offset") != null) {
                return Long.parseLong(response.getProperty("offset"));
            }
            log.warn("Failed to fetch offset for {} partition {}: {}. Defaulting to -1L.",
                    tp.getTopic(), tp.getPartitionId(), response.getErrorMessage());
            return -1L;
        });
    }

    // --- Push 模式 API (不支持) ---

    @Override
    public void subscribe(String topic, MessageListener listener) {
        throw new UnsupportedOperationException("subscribe is not supported in pull-style consumer. Please use pullResult() instead.");
    }

    @Override
    public void subscribe(String topic, BatchMessageListener listener) {
        throw new UnsupportedOperationException("subscribe is not supported in pull-style consumer. Please use pullResult() instead.");
    }

    @Override
    public void subscribe(String topic, OrderedMessageListener listener) {
        throw new UnsupportedOperationException("subscribe is not supported in pull-style consumer. Please use pullResult() instead.");
    }

    @Override
    public void subscribe(String topic, String filterExpression, MessageListener listener) {
        throw new UnsupportedOperationException("subscribe is not supported in pull-style consumer. Please use pullResult() instead.");
    }

    @Override
    public void subscribe(String topic, String filterExpression, BatchMessageListener listener) {
        throw new UnsupportedOperationException("subscribe is not supported in pull-style consumer. Please use pullResult() instead.");
    }

    @Override
    public void subscribe(String topic, String filterExpression, OrderedMessageListener listener) {
        throw new UnsupportedOperationException("subscribe is not supported in pull-style consumer. Please use pullResult() instead.");
    }
}